#include "hybrid_a_star.h"
#include <algorithm>

namespace HybridAStar {

std::vector<State> HybridAStar::plan(const State& start, const State& goal, 
                                   const OccupancyGrid& map) {
    using OpenList = std::priority_queue<State*, std::vector<State*>, NodeCompare>;
    OpenList open;
    std::unordered_map<GridKey, State*, GridHash> closed;

    State* start_node = new State(start);
    start_node->h = heuristic(*start_node, goal);
    open.push(start_node);

    while (!open.empty()) {
        State* current = open.top();
        open.pop();
        GridKey key = discretize(*current);

        if (closed.find(key) != closed.end()) continue;
        closed[key] = current;

        if (isGoalReached(*current, goal)) 
            return tracePath(current);

        for (State& next : generateSuccessors(*current, map)) {
            GridKey next_key = discretize(next);
            if (closed.count(next_key)) continue;

            next.h = heuristic(next, goal);
            open.push(new State(next));
        }
    }
    return {};
}

std::vector<State> HybridAStar::generateSuccessors(const State& current, 
                                                 const OccupancyGrid& map) {
    std::vector<State> successors;
    constexpr double steer_ratios[] = {-1.0, -0.5, 0.0, 0.5, 1.0};
    
    for (double ratio : steer_ratios) {
        for (int gear : {1, -1}) {
            State next;
            next.gear = gear;
            next.steer = ratio * params_.max_steer;
            
            const double step = gear * step_;
            const double delta_theta = step * tan(next.steer) / params_.wheelbase;
            next.theta = current.theta + delta_theta;
            
            next.x = current.x + step * cos(next.theta);
            next.y = current.y + step * sin(next.theta);
            next.parent = const_cast<State*>(&current);
            next.g = current.g + std::abs(step);

            if (!checkCollision(next, map))
                successors.push_back(next);
        }
    }
    return successors;
}

bool HybridAStar::checkCollision(const State& s, const OccupancyGrid& map) const {
    const auto corners = getVehicleCorners(s);
    for (const auto& pt : corners) {
        if (map.isOccupied(pt.first, pt.second)) return true;
    }
    return checkRectCollision(s, map);
}

std::array<std::pair<double, double>, 4> 
HybridAStar::getVehicleCorners(const State& s) const {
    const double cos_t = cos(s.theta);
    const double sin_t = sin(s.theta);
    const double front = params_.wheelbase + params_.front_overhang;
    const double rear = -params_.rear_overhang;
    const double half_w = params_.width / 2;

    return {{
        {s.x + front * cos_t - half_w * sin_t,
         s.y + front * sin_t + half_w * cos_t},
        {s.x + front * cos_t + half_w * sin_t,
         s.y + front * sin_t - half_w * cos_t},
        {s.x + rear * cos_t - half_w * sin_t,
         s.y + rear * sin_t + half_w * cos_t},
        {s.x + rear * cos_t + half_w * sin_t,
         s.y + rear * sin_t - half_w * cos_t}
    }};
}

bool HybridAStar::checkRectCollision(const State& s, const OccupancyGrid& map) const {
    const auto corners = getVehicleCorners(s);
    const double step = std::min(map_res_ * 1.5, params_.width/4.0);
    
    // 长边采样
    const int length_samples = std::ceil(params_.length / step);
    for (int i = 0; i <= length_samples; ++i) {
        const double ratio = i / double(length_samples);
        const double x = corners[0].first * (1 - ratio) + corners[2].first * ratio;
        const double y = corners[0].second * (1 - ratio) + corners[2].second * ratio;
        if (map.isOccupied(x, y)) return true;
    }
    
    // 宽边采样
    const int width_samples = std::ceil(params_.width / step);
    for (int i = 0; i <= width_samples; ++i) {
        const double ratio = i / double(width_samples);
        const double x = corners[0].first * (1 - ratio) + corners[1].first * ratio;
        const double y = corners[0].second * (1 - ratio) + corners[1].second * ratio;
        if (map.isOccupied(x, y)) return true;
    }
    
    return false;
}

HybridAStar::GridKey HybridAStar::discretize(const State& s) const {
    return {
        static_cast<int>(s.x / map_res_),
        static_cast<int>(s.y / map_res_),
        static_cast<int>(s.theta / theta_res_)
    };
}

double HybridAStar::heuristic(const State& s, const State& goal) const {
    const double dx = goal.x - s.x;
    const double dy = goal.y - s.y;
    const double dtheta = std::abs(fmod(s.theta - goal.theta, 2*M_PI));
    return std::hypot(dx, dy) + 0.2 * dtheta;
}

bool HybridAStar::isGoalReached(const State& s, const State& goal) const {
    const double pos_err = std::hypot(s.x - goal.x, s.y - goal.y);
    const double theta_err = std::abs(fmod(s.theta - goal.theta, 2*M_PI));
    return pos_err < map_res_*2 && theta_err < (15 * M_PI/180);
}

std::vector<State> HybridAStar::tracePath(State* node) const {
    std::vector<State> path;
    while (node) {
        path.push_back(*node);
        node = node->parent;
    }
    std::reverse(path.begin(), path.end());
    return path;
}
} // namespace HybridAStar